import torch
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim

batch_size = 32
# Prepare dataset
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,))])

train_dataset = datasets.MNIST(root='../dataset/mnist',
                               train=True,
                               transform=transform,
                               download=True)
test_dataset = datasets.MNIST(root='../dataset/mnist',
                              train=False,
                              transform=transform,
                              download=True)
train_dataloader = DataLoader(dataset=train_dataset,
                              batch_size=batch_size,
                              shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset,
                             batch_size=batch_size,
                             shuffle=False)


# design model by class
class MnistNet(torch.nn.Module):
    def __init__(self):
        super(MnistNet, self).__init__()
        self.linear_01 = torch.nn.Linear(784, 512)
        self.linear_02 = torch.nn.Linear(512, 256)
        self.linear_03 = torch.nn.Linear(256, 128)
        self.linear_04 = torch.nn.Linear(128, 64)
        self.linear_05 = torch.nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = F.relu(self.linear_01(x))
        x = F.relu(self.linear_02(x))
        x = F.relu(self.linear_03(x))
        x = F.relu(self.linear_04(x))
        return self.linear_05(x)


mnistModel = MnistNet()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(torch.cuda.is_available())
mnistModel.to(device)
# Construct loss and optimizer
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(mnistModel.parameters(), lr=0.01, momentum=0.5)


# Set train cycle
def train(epoch):
    mnistModel.train()
    running_loss = 0.0
    for batch_idx, data in enumerate(train_dataloader, 0):
        inputs, target = data
        # Set GPU
        inputs, target = inputs.to(device), target.to(device)
        # Forward + Backward + Update
        outputs = mnistModel(inputs)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        running_loss += loss.item()
        if batch_idx % 300 == 299:
            print('[%d,%5d] loss: %.6f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0


def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_dataloader:
            mnistModel.eval()
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            print(labels.shape[0])
            outputs = mnistModel(images)
            _, predicted = torch.max(outputs.data, dim=1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('Accuracy on test set : %3d %%' % (100 * correct / total))


if __name__ == "__main__":
    # print(next(iter(train_dataloader))[0].shape)
    # for name, parameter in mnistModel.named_parameters():
    #     print(name, parameter, parameter.size())
    for epoch in range(30):
        train(epoch)
        test()
